"""
重排模型评估结果可视化模块 - 支持带分数的新格式
"""
import os
from typing import Dict, List, Any, Optional, Union
import matplotlib as mpl

# 设置matplotlib后端为非交互式
mpl.use('Agg')  # 必须在导入pyplot之前设置

# 尝试导入可视化库，如果不可用则在使用时提供警告
try:
    import matplotlib.pyplot as plt
    import seaborn as sns
    import pandas as pd
    import numpy as np
    from matplotlib import font_manager
    
    # 设置中文字体支持 - 使用更可靠的方法
    def setup_chinese_font():
        # 尝试多种方法设置中文字体
        # 方法1: 使用系统中文字体
        font_dirs = []
        font_files = []
        
        # Windows常见中文字体路径
        windows_font_paths = [
            'C:/Windows/Fonts/simhei.ttf',  # 黑体
            'C:/Windows/Fonts/simsun.ttc',  # 宋体
            'C:/Windows/Fonts/msyh.ttc',    # 微软雅黑
            'C:/Windows/Fonts/simkai.ttf',  # 楷体
        ]
        
        # 检查字体是否存在并添加
        for font_path in windows_font_paths:
            if os.path.exists(font_path):
                font_dirs.append(os.path.dirname(font_path))
                font_files.append(font_path)
                break
        
        if font_files:
            # 注册找到的第一个中文字体
            font_path = font_files[0]
            prop = font_manager.FontProperties(fname=font_path)
            mpl.rcParams['font.family'] = 'sans-serif'
            mpl.rcParams['font.sans-serif'] = [prop.get_name()] + mpl.rcParams['font.sans-serif']
            return True
        
        # 方法2: 使用matplotlib内置字体
        try:
            # 尝试使用DejaVu Sans (matplotlib内置字体，支持基本中文)
            mpl.rcParams['font.family'] = ['sans-serif']
            mpl.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial Unicode MS'] + mpl.rcParams['font.sans-serif']
            mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
            return True
        except:
            pass
        
        # 方法3: 使用纯英文标签，避免中文
        return False
    
    # 设置中文字体
    chinese_font_available = setup_chinese_font()
    if not chinese_font_available:
        print("警告: 未找到合适的中文字体，将使用英文标签")
    
    VISUALIZATION_AVAILABLE = True
except ImportError:
    VISUALIZATION_AVAILABLE = False
    print("警告: matplotlib, seaborn 或 pandas 库未安装，可视化功能将不可用。")
    print("请运行 'pip install matplotlib seaborn pandas' 安装所需依赖。")

def plot_metric_comparison(model_results: Dict[str, Dict[str, float]], 
                          metric: str, 
                          title: Optional[str] = None, 
                          output_path: Optional[str] = None):
    """
    绘制两个模型在某个指标上的对比柱状图
    
    参数:
        model_results: 键为模型名称，值为该模型的评估结果
                      每个模型结果为字典，键为问题ID，值为指标得分(直接是float值)
        metric: 要比较的指标名称（如'ndcg@5', 'avg_relevance@5', 'hit_rate@5'）
        title: 图表标题，如果为None则自动生成
        output_path: 输出图像的路径，如果为None则只显示不保存
    """
    if not VISUALIZATION_AVAILABLE:
        print(f"无法绘制图表 '{metric}'，缺少必要的库。")
        return
        
    # 提取每个模型的指标平均值
    model_names = list(model_results.keys())
    avg_metrics = []
    
    for model_name in model_names:
        # 转换为纯数字列表，再计算平均值
        numeric_values = []
        for value in model_results[model_name].values():
            if isinstance(value, (int, float)):
                numeric_values.append(float(value))
        
        if numeric_values:
            avg_metrics.append(float(np.mean(numeric_values)))
        else:
            avg_metrics.append(0.0)
    
    # 创建DataFrame
    df = pd.DataFrame({
        'Model': model_names,  # 使用英文标签避免中文乱码
        'Value': avg_metrics
    })
    
    # 设置样式
    sns.set_style("whitegrid")
    plt.figure(figsize=(10, 6))
    
    # 绘制柱状图 - 修复警告
    # 使用seaborn的新API方式
    ax = sns.barplot(data=df, x='Model', y='Value', hue='Model', legend=False)
    
    # 添加数值标签
    for i, value in enumerate(avg_metrics):
        plt.text(i, value + 0.01, f'{value:.3f}', ha='center', va='bottom')
    
    # 设置标题和标签
    if title is None:
        title = f"Model Performance Comparison - {metric}"
    plt.title(title, fontsize=15)
    plt.xlabel('Model', fontsize=12)
    plt.ylabel(metric, fontsize=12)
    plt.ylim(0, max(avg_metrics) * 1.2)  # 设置Y轴范围
    
    # 保存图像
    if output_path:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
    
    # 关闭图形，不显示
    plt.close()

def plot_bad_case_analysis(bge_rankings: List[int], 
                          qwen_rankings: List[int], 
                          gold_scores: List[float],
                          question: str,
                          title: Optional[str] = None, 
                          output_path: Optional[str] = None):
    """
    为Bad Case绘制排名和得分对比图
    
    参数:
        bge_rankings: BGE模型的知识片段排名列表
        qwen_rankings: Qwen模型的知识片段排名列表
        gold_scores: 每个知识片段的黄金标准得分
        question: 问题文本
        title: 图表标题，如果为None则自动生成
        output_path: 输出图像的路径，如果为None则只显示不保存
    """
    if not VISUALIZATION_AVAILABLE:
        print(f"无法绘制Bad Case分析图 '{question[:30]}...'，缺少必要的库。")
        return
        
    # 设置样式
    sns.set_style("whitegrid")
    plt.figure(figsize=(14, 8))
    
    # 获取前5个知识片段的索引
    bge_top5 = bge_rankings[:5]
    qwen_top5 = qwen_rankings[:5]
    
    # 获取这些知识片段的得分
    bge_scores = [gold_scores[idx] for idx in bge_top5]
    qwen_scores = [gold_scores[idx] for idx in qwen_top5]
    
    # 创建数据
    positions = [f'Top{i+1}' for i in range(5)]
    
    # 使用matplotlib直接绘制分组柱状图
    x = np.arange(len(positions))  # 位置标签的位置
    width = 0.35  # 柱状图的宽度
    
    fig, ax = plt.subplots(figsize=(14, 8))
    rects1 = ax.bar(x - width/2, bge_scores, width, label='bge_rerank')
    rects2 = ax.bar(x + width/2, qwen_scores, width, label='qwen3_4b')
    
    # 添加一些文本和标签
    ax.set_ylabel('Gold Standard Relevance Score', fontsize=12)
    ax.set_xlabel('Position', fontsize=12)
    ax.set_title(title if title else f"Bad Case Analysis (ID: {question[:20]}...)" if len(question) > 20 else question, fontsize=15)
    ax.set_xticks(x)
    ax.set_xticklabels(positions)
    ax.legend(title='Model')
    
    # 添加数值标签
    def autolabel(rects):
        for rect in rects:
            height = rect.get_height()
            ax.annotate(f'{height:.1f}',
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3 点垂直偏移
                        textcoords="offset points",
                        ha='center', va='bottom')
    
    autolabel(rects1)
    autolabel(rects2)
    
    plt.ylim(0, 1.2)  # 设置Y轴范围为0-1
    
    # 保存图像
    if output_path:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
    
    # 关闭图形，不显示
    plt.close()

def plot_score_distribution(bge_scores: List[float], 
                           qwen_scores: List[float], 
                           gold_scores: List[float],
                           question_id: str,
                           title: Optional[str] = None, 
                           output_path: Optional[str] = None):
    """
    绘制Bad Case中分数分布对比的折线图
    
    参数:
        bge_scores: BGE模型的分数列表
        qwen_scores: Qwen模型的分数列表
        gold_scores: 黄金标准相关性得分列表
        question_id: 问题ID
        title: 图表标题，如果为None则自动生成
        output_path: 输出图像的路径，如果为None则只显示不保存
    """
    if not VISUALIZATION_AVAILABLE:
        print(f"无法绘制分数分布对比图 '{question_id}'，缺少必要的库。")
        return
    
    # 设置样式
    sns.set_style("whitegrid")
    plt.figure(figsize=(12, 8))
    
    # 定义相关性得分分组
    relevance_labels = ['0-0.2', '0.2-0.4', '0.4-0.6', '0.6-0.8', '0.8-1.0']
    
    # 创建分组数据
    bge_by_group = [[] for _ in range(5)]
    qwen_by_group = [[] for _ in range(5)]
    
    # 将分数按黄金标准得分分组
    for i, gold_score in enumerate(gold_scores):
        bin_idx = min(int(gold_score * 5), 4)  # 将0-1分数映射到0-4的索引
        bge_by_group[bin_idx].append(bge_scores[i])
        qwen_by_group[bin_idx].append(qwen_scores[i])
    
    # 计算每组的平均分数
    bge_avg = []
    qwen_avg = []
    gold_avg = []
    
    for i in range(5):
        if bge_by_group[i]:
            bge_avg.append(sum(bge_by_group[i]) / len(bge_by_group[i]))
            qwen_avg.append(sum(qwen_by_group[i]) / len(qwen_by_group[i]))
            # 黄金标准得分的平均值（每组的中点）
            gold_avg.append((i * 0.2) + 0.1)  # 0.1, 0.3, 0.5, 0.7, 0.9
        else:
            # 如果该组没有数据，使用0
            bge_avg.append(0)
            qwen_avg.append(0)
            gold_avg.append((i * 0.2) + 0.1)
    
    # 绘制折线图
    x = np.arange(len(relevance_labels))
    
    plt.plot(x, bge_avg, marker='o', linewidth=2, label='BGE Rerank')
    plt.plot(x, qwen_avg, marker='s', linewidth=2, label='Qwen3 4B')
    plt.plot(x, gold_avg, marker='x', linestyle='--', linewidth=1, label='Ideal Score')
    
    # 添加数据标签
    for i in range(len(x)):
        plt.text(x[i], bge_avg[i] + 0.02, f"{bge_avg[i]:.2f}", ha='center', va='bottom')
        plt.text(x[i], qwen_avg[i] - 0.02, f"{qwen_avg[i]:.2f}", ha='center', va='top')
    
    # 设置X轴标签
    plt.xticks(x, relevance_labels)
    
    # 设置标题和标签
    if title is None:
        title = f"Score Distribution Comparison - Question {question_id}"
    plt.title(title, fontsize=15)
    plt.xlabel('Gold Standard Relevance Score Range', fontsize=12)
    plt.ylabel('Average Model Score', fontsize=12)
    plt.ylim(0, 1.1)  # 设置Y轴范围
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(title='Model')
    
    # 保存图像
    if output_path:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
    
    # 关闭图形，不显示
    plt.close() 